package com.demo.newLSTM;

import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.GradientNormalization;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.NDArrayIndex;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * Java编写的LSTM短期负荷预测 加入 季节、节假日、负荷特性、气象因素、政策因素、经济因素、负荷管理措施等影响因素
 * 案例改为预测4个小时的
 */
public class ShortTermLoadForecasting {
    public static void main(String[] args) {
        // 准备训练数据
        List<DataSet> trainData = prepareTrainingData();

        // 创建LSTM网络模型
        int numInputs = 7; // 输入特征数量
        int numOutputs = 1; // 输出负荷数量
        int lstmLayerSize = 128; // LSTM层大小
        int numHiddenLayers = 2; // 隐藏层数量

        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(123)
                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
                .weightInit(WeightInit.XAVIER)
                .updater(new Adam(0.01))
                .gradientNormalization(GradientNormalization.ClipElementWiseAbsoluteValue)
                .gradientNormalizationThreshold(0.5)
                .list()
                .layer(0, new LSTM.Builder()
                        .nIn(numInputs)
                        .nOut(lstmLayerSize)
                        .activation(Activation.TANH)
                        .build())
                .layer(1, new RnnOutputLayer.Builder(LossFunctions.LossFunction.MSE)
                        .activation(Activation.IDENTITY)
                        .nIn(lstmLayerSize)
                        .nOut(numOutputs)
                        .build())
                .build();

        MultiLayerNetwork model = new MultiLayerNetwork(conf);
        model.init();

        // 训练模型
        int numEpochs = 100;
        for (int i = 0; i < numEpochs; i++) {
            for (DataSet data : trainData) {
                model.fit(data);
            }
        }

        // 准备输入数据进行预测
        INDArray input = prepareInputData();

        // 进行预测
        INDArray output = model.rnnTimeStep(input);

        // 打印预测结果
        System.out.println("负荷预测结果：");
        for (int i = 0; i < output.size(2); i++) {
            System.out.println("时间步 " + (i + 1) + ": " + output.getDouble(0, 0, i));
        }
    }

    private static List<DataSet> prepareTrainingData() {
        // 准备训练数据的代码

        //负荷数据
        List<Double> load = Arrays.asList(10.0, 12.0, 15.0, 20.0, 18.0, 14.0, 12.0, 10.0, 9.0, 8.0, 9.0, 11.0, 13.0, 15.0, 18.0, 20.0);
        List<Double> temperature = Arrays.asList(20.0, 22.0, 25.0, 28.0, 30.0, 28.0, 25.0, 22.0, 20.0, 18.0, 16.0, 18.0, 20.0, 22.0, 25.0, 28.0);
        List<Double> holiday = Arrays.asList(0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0);
        List<Double> weekday = Arrays.asList(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.0, 2.0);
        //季节数据
        List<Double> season = Arrays.asList(1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 1.0, 1.0, 1.0, 2.0);

        // 将数据转换为INDArray对象
        INDArray loadArray = Nd4j.create(load.stream().mapToDouble(Double::doubleValue).toArray(), new int[]{load.size(), 1});
        INDArray temperatureArray = Nd4j.create(temperature.stream().mapToDouble(Double::doubleValue).toArray(), new int[]{temperature.size(), 1});
        INDArray holidayArray = Nd4j.create(holiday.stream().mapToDouble(Double::doubleValue).toArray(), new int[]{holiday.size(), 1});
        INDArray weekdayArray = Nd4j.create(weekday.stream().mapToDouble(Double::doubleValue).toArray(), new int[]{weekday.size(), 1});
        INDArray seasonArray = Nd4j.create(season.stream().mapToDouble(Double::doubleValue).toArray(), new int[]{season.size(), 1});

        // 创建DataSet对象
        List<DataSet> dataSets = new ArrayList<>();
        for (int i = 24; i < load.size(); i++) { // 前24个小时用于预热，不参与训练
            INDArray features = Nd4j.concat(1,
                    loadArray.get(NDArrayIndex.interval(i - 24, i), NDArrayIndex.all()),
                    temperatureArray.get(NDArrayIndex.interval(i - 24, i), NDArrayIndex.all()),
                    holidayArray.get(NDArrayIndex.interval(i - 24, i), NDArrayIndex.all()),
                    weekdayArray.get(NDArrayIndex.interval(i - 24, i), NDArrayIndex.all()),
                    seasonArray.get(NDArrayIndex.interval(i - 24, i), NDArrayIndex.all())
            );
            INDArray label = loadArray.get(NDArrayIndex.interval(i + 1, i + 5), NDArrayIndex.all()); // 预测未来4个小时的负荷
            dataSets.add(new DataSet(features, label));
        }

        // 将DataSet对象打乱顺序并返回
        Collections.shuffle(dataSets);
        return dataSets;
    }

    private static INDArray prepareInputData(){
        // 假设已经从数据库或其他数据源中获取到了负荷数据、气象数据等
        double[] load = {10.0, 12.0, 15.0, 20.0, 18.0, 14.0, 12.0, 10.0, 9.0, 8.0, 9.0, 11.0, 13.0, 15.0, 18.0, 20.0};
        //温度
        double[] temperature = {20.0, 22.0, 25.0, 28.0, 30.0, 28.0, 25.0, 22.0, 20.0, 18.0, 16.0, 18.0, 20.0, 22.0, 25.0, 28.0};
        //假期
        double[] holiday = {0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0};
        double[] weekday = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.0, 2.0};
        //季节数据
        double[] season = {1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 1.0, 1.0, 1.0};
        //政策
        double[] policy = {1.2, 1.5, 1.3, 1.4, 1.6, 1.7, 1.8, 1.9, 2.0, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7};


        int numExamples = load.length - 4;
        int numInputs = 7;
        INDArray input = Nd4j.zeros(numExamples, numInputs, 4);
        for (int i = 0; i < numExamples; i++) {
            input.putScalar(new int[] {i, 0, 0}, load[i]);
            input.putScalar(new int[] {i, 1, 0}, temperature[i]);
            input.putScalar(new int[] {i, 2, 0}, holiday[i]);
            input.putScalar(new int[] {i, 3, 0}, weekday[i]);
            input.putScalar(new int[] {i, 4, 0}, season[i]);
            input.putScalar(new int[] {i, 5, 0}, policy[i]);
//            input.putScalar(new int[] {i, 6, 0}, economy[i]);
//            input.putScalar(new int[] {i, 7, 0}, management[i]);
            input.putScalar(new int[] {i, 8, 0}, load[i + 1]);
            input.putScalar(new int[] {i, 0, 1}, load[i + 2]);
            input.putScalar(new int[] {i, 1, 1}, temperature[i + 2]);
            input.putScalar(new int[] {i, 2, 1}, holiday[i + 2]);
            input.putScalar(new int[] {i, 3, 1}, weekday[i + 2]);
            input.putScalar(new int[] {i, 4, 1}, season[i + 2]);
            input.putScalar(new int[] {i, 5, 1}, policy[i + 2]);
//            input.putScalar(new int[] {i, 6, 1}, economy[i + 2]);
//            input.putScalar(new int[] {i, 7, 1}, management[i + 2]);
            input.putScalar(new int[] {i, 8, 1}, load[i + 3]);
            input.putScalar(new int[] {i, 0, 2}, load[i + 4]);
            input.putScalar(new int[] {i, 1, 2}, temperature[i + 4]);
            input.putScalar(new int[] {i, 2, 2}, holiday[i + 4]);
            input.putScalar(new int[] {i, 3, 2}, weekday[i + 4]);
            input.putScalar(new int[] {i, 4, 2}, season[i + 4]);
            input.putScalar(new int[] {i, 5, 2}, policy[i + 4]);
//            input.putScalar(new int[] {i, 6, 2}, economy[i + 4]);
//            input.putScalar(new int[] {i, 7, 2}, management[i + 4]);
            input.putScalar(new int[] {i, 8, 2}, load[i + 5]);
            input.putScalar(new int[] {i, 0, 3}, load[i + 6]);
            input.putScalar(new int[] {i, 1, 3}, temperature[i + 6]);
            input.putScalar(new int[] {i, 2, 3}, holiday[i + 6]);
            input.putScalar(new int[] {i, 3, 3}, weekday[i + 6]);
            input.putScalar(new int[] {i, 4, 3}, season[i + 6]);
            input.putScalar(new int[] {i, 5, 3}, policy[i + 6]);
//            input.putScalar(new int[] {i, 6, 3}, economy[i + 6]);
//            input.putScalar(new int[] {i, 7, 3}, management[i + 6]);
        }
        return input;
    }

    private static INDArray prepareInputData1() {
        List<Double> load = Arrays.asList(10.0, 12.0, 15.0, 20.0, 18.0, 14.0, 12.0, 10.0, 9.0, 8.0, 9.0, 11.0, 13.0, 15.0, 18.0, 20.0);
        List<Double> temperature = Arrays.asList(20.0, 22.0, 25.0, 28.0, 30.0, 28.0, 25.0, 22.0, 20.0, 18.0, 16.0, 18.0, 20.0, 22.0, 25.0, 28.0);
        List<Double> holiday = Arrays.asList(0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0);
        List<Double> weekday = Arrays.asList(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.0, 2.0);
        //季节数据
        List<Double> season = Arrays.asList(1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 1.0, 1.0, 1.0, 2.0);

        int numExamples = load.size() - 4;
        int numInputs = 7;
        INDArray input = Nd4j.zeros(numExamples, numInputs, 4);
        for (int i = 0; i < numExamples; i++) {
            input.putScalar(new int[] {i, 0, 0}, load.get(i));
            input.putScalar(new int[] {i, 5, 0}, load.get(i + 1) );
            input.putScalar(new int[] {i, 6, 0}, load.get(i + 2) );
            input.putScalar(new int[] {i, 0, 1}, load.get(i + 3) );
            input.putScalar(new int[] {i, 5, 1}, load.get(i + 4) );
            //input.putScalar(new int[] {i, 6, 1}, load.get(i + 5) );
            //input.putScalar(new int[] {i, 0, 2}, load.get(i + 6) );
            //input.putScalar(new int[] {i, 5, 2}, load.get(i + 7) );
            //input.putScalar(new int[] {i, 6, 2}, load.get(i + 8) );
            //input.putScalar(new int[] {i, 0, 3}, load.get(i + 9) );

            input.putScalar(new int[] {i, 1, 0}, temperature.get(i));
            input.putScalar(new int[] {i, 1, 1}, temperature.get(i + 3) );
            //input.putScalar(new int[] {i, 1, 2}, temperature.get(i + 6) );
            //input.putScalar(new int[] {i, 1, 3}, temperature.get(i + 9) );

            input.putScalar(new int[] {i, 2, 0}, holiday.get(i));
            input.putScalar(new int[] {i, 2, 1}, holiday.get(i + 3) );
            //input.putScalar(new int[] {i, 2, 2}, holiday.get(i + 6) );
            //input.putScalar(new int[] {i, 2, 3}, holiday.get(i + 9) );

            input.putScalar(new int[] {i, 3, 0}, weekday.get(i));
            input.putScalar(new int[] {i, 3, 1}, weekday.get(i + 3) );
            //input.putScalar(new int[] {i, 3, 2}, weekday.get(i + 6) );
            //input.putScalar(new int[] {i, 3, 3}, weekday.get(i + 9) );

            input.putScalar(new int[] {i, 4, 0}, season.get(i));
            input.putScalar(new int[] {i, 4, 1}, season.get(i + 3) );
            //input.putScalar(new int[] {i, 4, 2}, season.get(i + 6) );
            //input.putScalar(new int[] {i, 4, 3}, season.get(i + 9) );
        }
        return input;
    }
}
